# YOLOv4 common modules and utils

import math
from copy import copy
from pathlib import Path

import numpy as np
import pandas as pd
import requests
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image

from utils.datasets import letterbox
from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh
from utils.plots import color_list, plot_one_box
from utils.torch_utils import time_synchronized

from models.common import autopad, Concat


class Mish(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self,x):
        x = x * (torch.tanh(F.softplus(x)))
        return x


class BottleneckM(nn.Module):
    # bottleneck with ConvM
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
        super(BottleneckM, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = ConvM(c1, c_, 1, 1)
        self.cv2 = ConvM(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        y = self.cv2(self.cv1(x))
        return x + y if self.add else y


class ConvM(nn.Module):
    #  convolution with Mish activation
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(ConvM, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = Mish() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def fuseforward(self, x):
        return self.act(self.conv(x))


class ConvL(nn.Module):
    #  convolution with LeakyReLU activation
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(ConvL, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.LeakyReLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def fuseforward(self, x):
        return self.act(self.conv(x))


class C5(nn.Module):
    # CSP Bottleneck with 5 convolutions (yolov4)
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(C5, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv0 = ConvM(c1, c2, 3, 2)
        self.cv1 = ConvM(c2, c_, 1, 1)
        self.cv2 = ConvM(c2, c_, 1, 1)
        self.cv3 = ConvM(c_, c_, 1, 1)
        self.cv4 = ConvM(2 * c_, c2, 1, 1)
        self.m = nn.Sequential(*[BottleneckM(c_, c_, shortcut, g, e=1.0) for _ in range(n)])

    def forward(self, x):
        y0 = self.cv0(x)
        y1 = self.cv3(self.m(self.cv1(y0)))
        y2 = self.cv2(y0)
        return self.cv4(torch.cat((y1, y2), dim=1))


class C4(nn.Module):
    # CSP Bottleneck with 4 convolutions (yolov4-tiny)
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(C4, self).__init__()
        c_ = int(c1 * e)  # hidden channels
        self.cv0 = ConvL(c1, c1, 3, 1)
        self.cv1 = ConvL(c1, c_, 3, 1)
        self.cv2 = ConvL(c_, c_, 3, 1)
        self.cv3 = ConvL(c1, c1, 1, 1)
        self.p = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        y0 = self.cv0(x)
        y1 = self.cv1(y0)
        y2 = self.cv2(y1)
        y3 = self.cv3(torch.cat((y1,y2), dim=1))
        return self.p(torch.cat((y0, y3), dim=1))